import os
import sys
import cv2
import numpy as np
import threading
import time
from datetime import datetime
from PyQt5.QtWidgets import QApplication, QWidget, QVBoxLayout, QPushButton, QLabel, QFileDialog, QSystemTrayIcon, QMenu, QAction
from PyQt5.QtCore import Qt, QThread, pyqtSignal, pyqtSlot
from PyQt5.QtGui import QImage, QPixmap, QIcon
from plyer import notification
from PIL import Image, ImageDraw, ImageFont

class FaceRecognitionThread(QThread):
    update_frame = pyqtSignal(np.ndarray)
    face_detected = pyqtSignal(str)
    
    def __init__(self, known_face_features, known_face_names):
        super().__init__()
        self.known_face_features = known_face_features
        self.known_face_names = known_face_names
        self.running = True
        self.last_notification_time = {}  # 用于跟踪每个人的最后通知时间
        self.notification_cooldown = 60  # 通知冷却时间（秒）
        
        # 加载OpenCV的人脸检测器
        self.face_detector = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
        
        # 使用简单的欧氏距离比较代替LBPH识别器
        self.face_recognizer = None
        self.threshold = 50000000  # 人脸匹配阈值，调高以提高匹配成功率
    
    def run(self):
        video_capture = cv2.VideoCapture(0)
        
        if not video_capture.isOpened():
            print("无法打开摄像头")
            return
            
        while self.running:
            ret, frame = video_capture.read()
            
            if not ret:
                continue
                
            # 转换为灰度图像进行人脸检测
            gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            
            # 检测人脸
            faces = self.face_detector.detectMultiScale(
                gray,
                scaleFactor=1.1,
                minNeighbors=5,
                minSize=(30, 30)
            )
            
            # 处理每个检测到的人脸
            for (x, y, w, h) in faces:
                face_roi = gray[y:y+h, x:x+w]
                
                # 调整大小以匹配训练数据
                face_roi = cv2.resize(face_roi, (100, 100))
                
                # 应用直方图均衡化以提高特征质量，与已知人脸处理保持一致
                face_roi = cv2.equalizeHist(face_roi)
                
                name = "未知"
                
                # 如果有已知人脸，尝试识别
                if len(self.known_face_features) > 0:
                    try:
                        # 使用欧氏距离比较人脸特征
                        min_dist = float('inf')
                        best_match_idx = -1
                        
                        # 将当前人脸与所有已知人脸比较
                        for i, known_face in enumerate(self.known_face_features):
                            # 计算欧氏距离
                            dist = np.sum((face_roi.flatten() - known_face.flatten()) ** 2)
                            # print(f"与 {self.known_face_names[i]} 的距离: {dist}")
                            if dist < min_dist:
                                min_dist = dist
                                best_match_idx = i
                        
                        print(f"最小距离: {min_dist}, 阈值: {self.threshold}")
                        # 如果最小距离小于阈值，认为是已知人脸
                        if min_dist < self.threshold:
                            name = self.known_face_names[best_match_idx]
                            # print(f"识别到: {name}")
                            
                            # 检查是否需要发送通知（避免频繁通知）
                            current_time = time.time()
                            if name != "未知" and (name not in self.last_notification_time or 
                                               current_time - self.last_notification_time[name] > self.notification_cooldown):
                                self.face_detected.emit(name)
                                self.last_notification_time[name] = current_time
                    except Exception as e:
                        print(f"识别错误: {e}")
                
                # 绘制人脸框
                cv2.rectangle(frame, (x, y), (x+w, y+h), (0, 0, 255), 2)
                
                # 绘制名称标签
                cv2.rectangle(frame, (x, y+h-35), (x+w, y+h), (0, 0, 255), cv2.FILLED)
                
                # 使用PIL绘制中文
                frame_pil = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
                draw = ImageDraw.Draw(frame_pil)
                
                # 尝试加载中文字体，如果失败则使用默认字体
                try:
                    # 使用系统字体
                    fontpath = "C:/Windows/Fonts/simhei.ttf"  # 黑体
                    font = ImageFont.truetype(fontpath, 24)
                except IOError:
                    # 如果找不到指定字体，使用默认字体
                    font = ImageFont.load_default()
                
                # 绘制文本
                draw.text((x+6, y+h-30), name, font=font, fill=(255, 255, 255))
                
                # 转换回OpenCV格式
                frame = cv2.cvtColor(np.array(frame_pil), cv2.COLOR_RGB2BGR)
            
            # 发送处理后的帧
            self.update_frame.emit(frame)
            
            # 短暂休眠以减少CPU使用
            time.sleep(0.01)
        
        video_capture.release()
    
    def stop(self):
        self.running = False
        self.wait()

class FaceRecognitionApp(QWidget):
    def __init__(self):
        super().__init__()
        self.known_face_features = []
        self.known_face_names = []
        self.face_thread = None
        self.tray_icon = None
        self.images_folder = os.path.join(os.path.dirname(os.path.abspath(__file__)), "known_faces")
        
        # 确保known_faces文件夹存在
        if not os.path.exists(self.images_folder):
            os.makedirs(self.images_folder)
            
        self.initUI()
        self.loadKnownFaces()
        
    def initUI(self):
        self.setWindowTitle('人脸识别应用')
        self.setGeometry(100, 100, 640, 520)
        
        layout = QVBoxLayout()
        
        # 视频显示区域
        self.image_label = QLabel()
        self.image_label.setAlignment(Qt.AlignCenter)
        layout.addWidget(self.image_label)
        
        # 人脸库由用户手动维护，不需要添加人脸按钮
        
        # 开始/停止按钮
        self.toggle_btn = QPushButton('开始识别')
        self.toggle_btn.clicked.connect(self.toggleRecognition)
        layout.addWidget(self.toggle_btn)
        
        # 最小化到托盘按钮
        self.minimize_btn = QPushButton('最小化到托盘')
        self.minimize_btn.clicked.connect(self.hide)
        layout.addWidget(self.minimize_btn)
        
        self.setLayout(layout)
        
        # 创建系统托盘图标
        self.createTrayIcon()
        
    def createTrayIcon(self):
        # 创建托盘图标
        self.tray_icon = QSystemTrayIcon(self)
        self.tray_icon.setIcon(QIcon(self.style().standardIcon(QApplication.style().SP_ComputerIcon)))
        
        # 创建托盘菜单
        tray_menu = QMenu()
        
        show_action = QAction("显示", self)
        show_action.triggered.connect(self.show)
        tray_menu.addAction(show_action)
        
        hide_action = QAction("隐藏", self)
        hide_action.triggered.connect(self.hide)
        tray_menu.addAction(hide_action)
        
        quit_action = QAction("退出", self)
        quit_action.triggered.connect(self.closeApp)
        tray_menu.addAction(quit_action)
        
        self.tray_icon.setContextMenu(tray_menu)
        self.tray_icon.activated.connect(self.trayIconActivated)
        self.tray_icon.show()
    
    def trayIconActivated(self, reason):
        if reason == QSystemTrayIcon.DoubleClick:
            if self.isHidden():
                self.show()
            else:
                self.hide()
    
    def closeEvent(self, event):
        # 隐藏窗口而不是关闭应用
        event.ignore()
        self.hide()
        
    def closeApp(self):
        # 停止人脸识别线程
        if self.face_thread and self.face_thread.isRunning():
            self.face_thread.stop()
            self.face_thread.wait()
        QApplication.quit()
        
    def loadKnownFaces(self):
        # 清空现有的人脸数据
        self.known_face_features = []
        self.known_face_names = []
        
        # 从known_faces文件夹加载人脸
        if os.path.exists(self.images_folder):
            # 检查是否有子文件夹（每个人一个文件夹）
            has_person_folders = False
            for item in os.listdir(self.images_folder):
                if os.path.isdir(os.path.join(self.images_folder, item)):
                    has_person_folders = True
                    break
            
            if has_person_folders:
                # 按人名文件夹组织的结构
                for person_name in os.listdir(self.images_folder):
                    person_dir = os.path.join(self.images_folder, person_name)
                    if os.path.isdir(person_dir):
                        # 处理该人物文件夹中的所有图片
                        for image_name in os.listdir(person_dir):
                            if image_name.lower().endswith(('.jpg', '.jpeg', '.png')):
                                image_path = os.path.join(person_dir, image_name)
                                self.processKnownFace(image_path, person_name)
            else:
                # 直接存放图片的结构
                for filename in os.listdir(self.images_folder):
                    if filename.lower().endswith(('.jpg', '.jpeg', '.png')):
                        image_path = os.path.join(self.images_folder, filename)
                        # 使用文件名作为人名（去掉扩展名）
                        name = os.path.splitext(filename)[0]
                        self.processKnownFace(image_path, name)
    
    def processKnownFace(self, image_path, name):
        # 加载图片并检测人脸
        image = cv2.imread(image_path)
        if image is None:
            print(f"警告: 无法加载图像 {image_path}")
            return
        
        # 转换为灰度图
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        
        # 检测人脸
        face_detector = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
        faces = face_detector.detectMultiScale(gray, 1.1, 5, minSize=(30, 30))
        
        if len(faces) > 0:
            # 获取第一个人脸
            x, y, w, h = faces[0]
            face_roi = gray[y:y+h, x:x+w]
            
            # 调整大小为统一尺寸
            face_roi = cv2.resize(face_roi, (100, 100))
            
            # 应用直方图均衡化以提高特征质量
            face_roi = cv2.equalizeHist(face_roi)
            
            # 添加到已知人脸列表
            self.known_face_features.append(face_roi)
            self.known_face_names.append(name)
            print(f"已加载人脸: {name}, 路径: {image_path}")
    
    def toggleRecognition(self):
        if self.face_thread and self.face_thread.isRunning():
            self.stopRecognition()
            self.toggle_btn.setText('开始识别')
        else:
            self.startRecognition()
            self.toggle_btn.setText('停止识别')
    
    def startRecognition(self):
        if not self.face_thread or not self.face_thread.isRunning():
            self.face_thread = FaceRecognitionThread(self.known_face_features, self.known_face_names)
            self.face_thread.update_frame.connect(self.updateFrame)
            self.face_thread.face_detected.connect(self.onFaceDetected)
            self.face_thread.start()
    
    def stopRecognition(self):
        if self.face_thread and self.face_thread.isRunning():
            self.face_thread.stop()
    
    @pyqtSlot(np.ndarray)
    def updateFrame(self, frame):
        # 将OpenCV的BGR格式转换为RGB
        rgb_image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        h, w, ch = rgb_image.shape
        bytes_per_line = ch * w
        
        # 创建QImage
        qt_image = QImage(rgb_image.data, w, h, bytes_per_line, QImage.Format_RGB888)
        
        # 创建QPixmap并显示
        pixmap = QPixmap.fromImage(qt_image)
        self.image_label.setPixmap(pixmap.scaled(self.image_label.width(), self.image_label.height(), Qt.KeepAspectRatio))
    
    @pyqtSlot(str)
    def onFaceDetected(self, name):
        # 发送桌面通知
        notification.notify(
            title='人脸识别警报',
            message=f'检测到已知人物: {name}',
            app_name='人脸识别系统',
            app_icon=None,  # 可以设置自定义图标路径
            timeout=10
        )
        
        # 如果窗口隐藏，显示托盘通知
        if not self.isVisible() and self.tray_icon:
            self.tray_icon.showMessage(
                '人脸识别警报',
                f'检测到已知人物: {name}',
                QSystemTrayIcon.Information,
                5000
            )

def main():
    app = QApplication(sys.argv)
    window = FaceRecognitionApp()
    window.show()
    sys.exit(app.exec_())

if __name__ == '__main__':
    main()